import os
import random
import moviepy.editor as mpy
import glob
import argparse


def check_path(path):
    # video_path
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)


def reset_fps(video):
    if args.fps != 25:
        video = video.set_fps(args.fps)
        print('reset video fps', args.fps)
    else:
        pass
    return video


def generate_video_by_way1(video_path):
    file_name = os.path.basename(video_path)
    # Load the video
    video = mpy.VideoFileClip(video_path)
    video = reset_fps(video)

    # Get total frame number
    total_frames = video.reader.nframes

    try:
        # Check if total frames are more than 4 (we need at least 4 frames to make a swap)
        if total_frames > 4:
            # Generate random frame indices for swapping
            swap_frames = sorted(random.sample(range(total_frames), 4))
            # Extract the clips to be swapped
            clip1 = video.subclip(swap_frames[0] / video.fps, swap_frames[1] / video.fps)
            clip2 = video.subclip(swap_frames[2] / video.fps, swap_frames[3] / video.fps)

            # Create the new video by concatenating the parts
            new_video = mpy.concatenate_videoclips([video.subclip(0, swap_frames[0] / video.fps),
                                                    clip2,
                                                    video.subclip(swap_frames[1] / video.fps,
                                                                  swap_frames[2] / video.fps),
                                                    clip1,
                                                    video.subclip(swap_frames[3] / video.fps)])

            new_video_name = file_name.split('.mp4')[0]
            for swap_index in swap_frames:
                new_video_name = new_video_name + '_' + str(swap_index)
            new_video_name += '_exchange.mp4'
            print(new_video_name)
            # Export the new video
            new_video.write_videofile(args.forge_video_path + '/' + new_video_name)

        else:
            print("Total frames are less than 4. Cannot perform swap operation.")
    except Exception as e:
        print(str(e))


def generate_video_by_way2(video_path):
    file_name = os.path.basename(video_path)
    # Load the video
    video = mpy.VideoFileClip(video_path)

    video = reset_fps(video)

    # Get total frame number
    total_frames = video.reader.nframes

    try:
        # Check if total frames are more than 2 (we need at least 2 frames to make a copy)
        if total_frames > 2:
            # Generate random frame index for inserting the copied frames
            insert_frame_index = random.randint(1, total_frames - 2)
            # Generate random frame index for copying (ensure it's within valid range)
            copy_frame_index = random.randint(0, total_frames - 3)
            print(total_frames, insert_frame_index, copy_frame_index)
            # Extract the frame to be copied
            frame = video.get_frame(copy_frame_index / video.fps)

            # Create the new video by inserting the copied frame at the specified index
            new_video = mpy.concatenate_videoclips([video.subclip(0, insert_frame_index / video.fps),
                                                    mpy.ImageClip(frame).set_duration(1 / video.fps),
                                                    mpy.ImageClip(frame).set_duration(1 / video.fps),
                                                    video.subclip(insert_frame_index / video.fps)])

            new_video_name = file_name.split('.mp4')[0]
            new_video_name = new_video_name + '_' + str(insert_frame_index + 1) + '_' + str(
                insert_frame_index + 2) + '_copy.mp4'

            # Export the new video
            new_video.write_videofile(args.forge_video_path + '/' + new_video_name)

        else:
            print("Total frames are less than 2. Cannot perform copying operation.")
    except Exception as e:
        print(str(e))


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='This is an example of command-line arguments')

    # Add command-line arguments
    parser.add_argument('--folder_path', type=str,
                        default="E:/code/Video Frame Interpolation Forgery Detection/VIFFD - A dataset for detecting video inter-frame forgeries/Training Data")
    parser.add_argument('--fps', type=int, default=25)
    parser.add_argument('--forge_video_path', type=str, default='forge_video_data')

    # Parse the command-line arguments
    args = parser.parse_args()
    check_path(args.forge_video_path)

    mp4_files = glob.glob(args.folder_path + "/*.mp4")

    # Shuffle the file list randomly
    random.shuffle(mp4_files)

    # Calculate the count of half the files
    half_count = len(mp4_files) // 2

    # Randomly select half the files to execute generate_video_by_way1
    for file_path in mp4_files[:half_count]:
        generate_video_by_way1(file_path)

    # Execute generate_video_by_way2 for the remaining half of the files
    for file_path in mp4_files[half_count:]:
        generate_video_by_way2(file_path)
